import os, glob

try:
    import numpy as np
except:
    raise ImportError('The numpy library cannot be found!')

try:
    from cdar.models.Shallow_Water_2D.ShallowWater2DModel import ShallowWaterModelClass as model
except:
    raise ImportError('The model library cannot be found!')
        
try:
    import cdar.models.Shallow_Water_2D.ShallowWater2DModelVis as vis
except:
    raise ImportError('The vis library cannot be found!')

try:
    from cdar.utilities import namelist_read as nml
except:
    raise ImportError('The nml library cannot be found!')

DataFilePath='./data'  # Path to store or read FG, B and OBS

# Read in parameters for model configurations
params = nml.namelist_read(DataFilePath)

# Set up the model    
truth = model(params)  

# Print domain info 
print truth  

# Observation Errors    
obs_u_error = 0.002  # Units: m/s
obs_v_error = 0.002  # Units: m/s
obs_h_error = 0.031  # Units: m

# Define the frequency to sample the data
obs_time = [time for time in truth.history if time % int(params['history_interval']) == 0 ]

# Define how many observations for each variables on each time
obs_num_u = [200]*len(obs_time)
obs_num_v = [200]*len(obs_time)
obs_num_h = [200]*len(obs_time)

#Sanity check

if not len(obs_time) ==  len(obs_num_u) or not len(obs_num_v) ==  len(obs_num_u) or \
   not len(obs_num_v) == len(obs_num_h) or not len(obs_time) == len(obs_num_h) :
    raise Exception (" obs_time and obs_num_u,v,h should have the same length")

# Plot the 0H field
#vis.plot_3d(truth.x, truth.y, truth.h, 'H at Time ='+str(0))

# Propogate the model
output = truth.model_integration()

# Generate the psedo observations on the shape w/omthe boundarie    
obs_u = np.empty_like( truth.u )
obs_v = np.empty_like( truth.v )
obs_h = np.empty_like( truth.h )

obs_t_u = np.empty_like( truth.u )
obs_t_v = np.empty_like( truth.v )
obs_t_h = np.empty_like( truth.h )

#Remove all the pre-existed OBS, OBS_TRUTH and TRUTH files
for filename in glob.glob(os.path.join(DataFilePath,'OBS_'+truth.Name+'_*')):
    os.remove(filename) 

for filename in glob.glob(os.path.join(DataFilePath,'OBS_TRUTH_'+truth.Name+'_*')):
    os.remove(filename) 
    
for filename in glob.glob(os.path.join(DataFilePath,'TRUTH_'+truth.Name+'_*')):
    os.remove(filename) 
    
(obs_u, obs_v, obs_h) = np.split(truth.strip_off_boundaries(np.vstack((obs_u,obs_v,obs_h))), truth.VarDim)
(obs_t_u, obs_t_v, obs_t_h) = np.split(truth.strip_off_boundaries(np.vstack((obs_t_u,obs_t_v,obs_t_h))), truth.VarDim)

# Generate pseudo observations.
for time in obs_time :
    
    print "+"*100    
    print 'Generating observations on time= %d from output' % time
    u,v,h = np.split( truth.strip_off_boundaries( output[np.where(truth.history==time),:].reshape(truth.ic.shape) ),truth.VarDim )

    obs_u.fill(np.nan)
    obs_v.fill(np.nan)
    obs_h.fill(np.nan)
    
    # Used to save the truth observation
    obs_t_u.fill(np.nan)
    obs_t_v.fill(np.nan)
    obs_t_h.fill(np.nan)
             
    index = obs_time.index(time)
    # For U variables
    # Randomly pick up obs_num point to make pseudo observations
    obs_x = np.floor(truth.nx * np.random.sample(obs_num_u[index])).astype(int)
    obs_y = np.floor(truth.ny * np.random.sample(obs_num_u[index])).astype(int)
        
    # Add random noise on the truth to make the observations            
    for i,j in zip(obs_x, obs_y) :
        obs_u[i,j] = u[i,j] + obs_u_error*np.random.randn(1)
        obs_t_u[i,j] = u[i,j]
        
    # For V variables
    # Randomly pick up obs_num point to make psedo observations
    obs_x = np.floor(truth.nx * np.random.sample(obs_num_v[index])).astype(int)
    obs_y = np.floor(truth.ny * np.random.sample(obs_num_v[index])).astype(int)
        
    # Add random noise on the truth to make the observations        
    for i,j in zip(obs_x, obs_y) :
        obs_v[i,j] = v[i,j] + obs_v_error*np.random.randn(1)
        obs_t_v[i,j] = v[i,j]
        
    # For H variables
    # Randomly pick up obs_num point to make psedo observations
    obs_x = np.floor(truth.nx * np.random.sample(obs_num_h[index])).astype(int)
    obs_y = np.floor(truth.ny * np.random.sample(obs_num_h[index])).astype(int)
        
    # Add random noise on the truth to make the observations        
    for i,j in zip(obs_x, obs_y) :
        obs_h[i,j] = h[i,j] + obs_h_error*np.random.randn(1)
        obs_t_h[i,j] = h[i,j]
    
    vis.plot_2d(truth.x, truth.y, obs_h, 'H at Time ='+str(time))    
    print "Saving observations on %d to file %s" % (time, os.path.join(DataFilePath,'OBS_'+truth.Name+'_'+str(int(time)))+'.npz')
    np.savez(os.path.join(DataFilePath,'OBS_'+truth.Name+'_'+str(int(time))), OBS=np.vstack((obs_u,obs_v,obs_h)), \
                      OBS_ERROR=np.vstack((obs_u_error,obs_v_error,obs_h_error)))
     
    print "Saving truth observations on %d to file %s" % (time, os.path.join(DataFilePath,'OBS_TRUTH_'+truth.Name+'_'+str(int(time)))+'.npz')
    np.savez(os.path.join(DataFilePath,'OBS_TRUTH_'+truth.Name+'_'+str(int(time))), OBS=np.vstack((obs_t_u,obs_t_v,obs_t_h)))  
                       
    # Plot and save the truth at time 
    print "Saving truth on %d to file %s" % (time, os.path.join(DataFilePath,'TRUTH_'+truth.Name+'_'+str(int(time)))+'.npz')
    np.savez(os.path.join(DataFilePath,'TRUTH_'+truth.Name+'_'+str(int(time))), TRUTH=output[np.where(truth.history == time),:].reshape(truth.ic.shape))

    # Verifiing
    diff = np.vstack((obs_t_u,obs_t_v,obs_t_h)) - truth.strip_off_boundaries( output[np.where(truth.history == time),:].reshape(truth.ic.shape) )
    diff[np.isnan(diff)] = 0.0
    
    print np.sum(diff)